use base64::{engine::general_purpose::STANDARD, Engine};

const KEY: [u8; 16] = [0u8; 16];
const IV: [u8; 16] = [1u8; 16];

const D: [u32; 16] = [
    0x44D7, 0x26BC, 0x626B, 0x135E, 0x5789, 0x35E2, 0x7135, 0x09AF,
    0x4D78, 0x2F13, 0x6BC4, 0x1AF1, 0x5E26, 0x3C4D, 0x789A, 0x47AC,
];

const S0: [u8; 256] = [
    0x3e,0x72,0x5b,0x47,0xca,0xe0,0x00,0x33,0x04,0xd1,0x54,0x98,0x09,0xb9,0x6d,0xcb, 
    0x7b,0x1b,0xf9,0x32,0xaf,0x9d,0x6a,0xa5,0xb8,0x2d,0xfc,0x1d,0x08,0x53,0x03,0x90, 
    0x4d,0x4e,0x84,0x99,0xe4,0xce,0xd9,0x91,0xdd,0xb6,0x85,0x48,0x8b,0x29,0x6e,0xac, 
    0xcd,0xc1,0xf8,0x1e,0x73,0x43,0x69,0xc6,0xb5,0xbd,0xfd,0x39,0x63,0x20,0xd4,0x38, 
    0x76,0x7d,0xb2,0xa7,0xcf,0xed,0x57,0xc5,0xf3,0x2c,0xbb,0x14,0x21,0x06,0x55,0x9b, 
    0xe3,0xef,0x5e,0x31,0x4f,0x7f,0x5a,0xa4,0x0d,0x82,0x51,0x49,0x5f,0xba,0x58,0x1c, 
    0x4a,0x16,0xd5,0x17,0xa8,0x92,0x24,0x1f,0x8c,0xff,0xd8,0xae,0x2e,0x01,0xd3,0xad, 
    0x3b,0x4b,0xda,0x46,0xeb,0xc9,0xde,0x9a,0x8f,0x87,0xd7,0x3a,0x80,0x6f,0x2f,0xc8, 
    0xb1,0xb4,0x37,0xf7,0x0a,0x22,0x13,0x28,0x7c,0xcc,0x3c,0x89,0xc7,0xc3,0x96,0x56, 
    0x07,0xbf,0x7e,0xf0,0x0b,0x2b,0x97,0x52,0x35,0x41,0x79,0x61,0xa6,0x4c,0x10,0xfe, 
    0xbc,0x26,0x95,0x88,0x8a,0xb0,0xa3,0xfb,0xc0,0x18,0x94,0xf2,0xe1,0xe5,0xe9,0x5d, 
    0xd0,0xdc,0x11,0x66,0x64,0x5c,0xec,0x59,0x42,0x75,0x12,0xf5,0x74,0x9c,0xaa,0x23, 
    0x0e,0x86,0xab,0xbe,0x2a,0x02,0xe7,0x67,0xe6,0x44,0xa2,0x6c,0xc2,0x93,0x9f,0xf1, 
    0xf6,0xfa,0x36,0xd2,0x50,0x68,0x9e,0x62,0x71,0x15,0x3d,0xd6,0x40,0xc4,0xe2,0x0f, 
    0x8e,0x83,0x77,0x6b,0x25,0x05,0x3f,0x0c,0x30,0xea,0x70,0xb7,0xa1,0xe8,0xa9,0x65, 
    0x8d,0x27,0x1a,0xdb,0x81,0xb3,0xa0,0xf4,0x45,0x7a,0x19,0xdf,0xee,0x78,0x34,0x60 
];

const S1: [u8;256] = [
    0x55,0xc2,0x63,0x71,0x3b,0xc8,0x47,0x86,0x9f,0x3c,0xda,0x5b,0x29,0xaa,0xfd,0x77, 
    0x8c,0xc5,0x94,0x0c,0xa6,0x1a,0x13,0x00,0xe3,0xa8,0x16,0x72,0x40,0xf9,0xf8,0x42, 
    0x44,0x26,0x68,0x96,0x81,0xd9,0x45,0x3e,0x10,0x76,0xc6,0xa7,0x8b,0x39,0x43,0xe1, 
    0x3a,0xb5,0x56,0x2a,0xc0,0x6d,0xb3,0x05,0x22,0x66,0xbf,0xdc,0x0b,0xfa,0x62,0x48, 
    0xdd,0x20,0x11,0x06,0x36,0xc9,0xc1,0xcf,0xf6,0x27,0x52,0xbb,0x69,0xf5,0xd4,0x87, 
    0x7f,0x84,0x4c,0xd2,0x9c,0x57,0xa4,0xbc,0x4f,0x9a,0xdf,0xfe,0xd6,0x8d,0x7a,0xeb, 
    0x2b,0x53,0xd8,0x5c,0xa1,0x14,0x17,0xfb,0x23,0xd5,0x7d,0x30,0x67,0x73,0x08,0x09, 
    0xee,0xb7,0x70,0x3f,0x61,0xb2,0x19,0x8e,0x4e,0xe5,0x4b,0x93,0x8f,0x5d,0xdb,0xa9, 
    0xad,0xf1,0xae,0x2e,0xcb,0x0d,0xfc,0xf4,0x2d,0x46,0x6e,0x1d,0x97,0xe8,0xd1,0xe9, 
    0x4d,0x37,0xa5,0x75,0x5e,0x83,0x9e,0xab,0x82,0x9d,0xb9,0x1c,0xe0,0xcd,0x49,0x89, 
    0x01,0xb6,0xbd,0x58,0x24,0xa2,0x5f,0x38,0x78,0x99,0x15,0x90,0x50,0xb8,0x95,0xe4, 
    0xd0,0x91,0xc7,0xce,0xed,0x0f,0xb4,0x6f,0xa0,0xcc,0xf0,0x02,0x4a,0x79,0xc3,0xde, 
    0xa3,0xef,0xea,0x51,0xe6,0x6b,0x18,0xec,0x1b,0x2c,0x80,0xf7,0x74,0xe7,0xff,0x21, 
    0x5a,0x6a,0x54,0x1e,0x41,0x31,0x92,0x35,0xc4,0x33,0x07,0x0a,0xba,0x7e,0x0e,0x34, 
    0x88,0xb1,0x98,0x7c,0xf3,0x3d,0x60,0x6c,0x7b,0xca,0xd3,0x1f,0x32,0x65,0x04,0x28, 
    0x64,0xbe,0x85,0x9b,0x2f,0x59,0x8a,0xd7,0xb0,0x25,0xac,0xaf,0x12,0x03,0xe2,0xf2 
];

struct ZucState {
    lfsr: [u32; 16], // 16个31bit寄存器单元，具有两种工作模式：初始化模式和工作模式
    r1: u32,         // 32位寄存器R1
    r2: u32,         // 32位寄存器R2
}

impl ZucState {
    pub fn new(key: &[u8], iv: &[u8]) -> Self {
        let mut state = ZucState {
            lfsr: [0; 16],
            r1: 0,
            r2: 0,
        };
        // Key and IV loading
        for i in 0..16 {
            state.lfsr[i] = ((key[i] as u32) << 23) | ((D[i] as u32) << 8) | (iv[i] as u32);
        }
        // Perform initialization mode
        for _ in 0..32 {
            let x = state.bit_reorganization();
            let w = state.f_function(x);
            state.lfsr_with_initialization_mode(w>>1);
        }
        state
    }

    fn mul_by_pow2(x: u32, k: u32) -> u32 {
        ((x << k) | (x >> (31 - k))) & 0x7FFFFFFF
    }

    fn add_mod31(a: u32, b: u32) -> u32 {
        let c = a.wrapping_add(b);
        (c & 0x7FFFFFFF) + (c >> 31)
    }

    fn lfsr_with_initialization_mode(&mut self, u: u32) {
        let mut f = self.lfsr[0];

        for (s, k) in [8, 20, 21, 17, 15].iter().zip([0, 4, 10, 13, 15].iter()) {
            let v = Self::mul_by_pow2(self.lfsr[*k], *s);
            f = Self::add_mod31(f, v);
        }

        f = Self::add_mod31(f, u);
        for i in 0..15 {
            self.lfsr[i] = self.lfsr[i+1];
        }
        self.lfsr[15] = f;
    }

    fn lfsr_with_work_mode(&mut self) {
        let mut f = self.lfsr[0];

        for (s, k) in [8, 20, 21, 17, 15].iter().zip([0, 4, 10, 13, 15].iter()) {
            let v = Self::mul_by_pow2(self.lfsr[*k], *s);
            f = Self::add_mod31(f, v);
        }

        for i in 0..15 {
            self.lfsr[i] = self.lfsr[i+1];
        }
        self.lfsr[15] = f;
    }

    fn bit_reorganization(&mut self) -> [u32; 4] {
        [
            ((self.lfsr[15] & 0x7FFF8000) << 1) | (self.lfsr[14] & 0xFFFF),
            ((self.lfsr[11] & 0xFFFF) << 16) | (self.lfsr[9] >> 15),
            ((self.lfsr[7] & 0xFFFF) << 16) | (self.lfsr[5] >> 15),
            ((self.lfsr[2] & 0xFFFF) << 16) | (self.lfsr[0] >> 15),
        ]
    }

    fn f_function(&mut self, x: [u32; 4]) -> u32 {
        let w = (x[0] ^ self.r1).wrapping_add(self.r2);
        let w1 = self.r1.wrapping_add(x[1]);
        let w2 = self.r2 ^ x[2];

        let u = Self::l1((w1<<16) | (w2>>16));
        let v = Self::l2((w2<<16) | (w1>>16));

        self.r1 = Self::make_u32(S0[((u>>24) as u8) as usize], S1[((u>>16) as u8) as usize], S0[((u>>8) as u8) as usize], S1[(u as u8) as usize]);
        self.r2 = Self::make_u32(S0[((v>>24) as u8) as usize], S1[((v>>16) as u8) as usize], S0[((v>>8) as u8) as usize], S1[(v as u8) as usize]);

        w
    }

    fn l1(x: u32) -> u32 {
        x ^ (x.rotate_left(2)) ^ (x.rotate_left(10)) ^ (x.rotate_left(18)) ^ (x.rotate_left(24))
    }

    fn l2(x: u32) -> u32 {
        x ^ (x.rotate_left(8)) ^ (x.rotate_left(14)) ^ (x.rotate_left(22)) ^ (x.rotate_left(30))
    }

    fn make_u32(a: u8, b: u8, c: u8, d: u8) -> u32 {
        ((a as u32) << 24) | ((b as u32) << 16) | ((c as u32) << 8) | (d as u32)
    }

    pub fn generate_keystream(&mut self, length: usize) -> Vec<u8> {
        let mut keystream: Vec<u8> = Vec::with_capacity(length);
        let x = self.bit_reorganization();
        let _ = self.f_function(x);
        self.lfsr_with_work_mode();

        for _ in 0..length {
            let x = self.bit_reorganization();
            let z = self.f_function(x) ^ x[3];
            self.lfsr_with_work_mode();
            keystream.push((z >> 24) as u8);
            keystream.push((z >> 16) as u8);
            keystream.push((z >> 8) as u8);
            keystream.push(z as u8);
        }
        keystream
    }
}

fn pkcs7_padding(data: &[u8], block_size: usize)-> Vec<u8> {
    let padding = block_size - data.len() % block_size;
    let mut result = Vec::from(data);
    result.extend(vec![padding as u8; padding]);
    result
}

pub fn encryption(input: String) -> String {
    let mut data = pkcs7_padding(input.as_bytes(), 4);
    let mut state = ZucState::new(&KEY, &IV);
    let keystream = state.generate_keystream(data.len()/4);
    for i in 0..data.len() {
        data[i] ^= keystream[i];
    }
    STANDARD.encode(&data)
}
